#!/bin/bash
# 清理所有 Ray worker 进程的脚本
# 用于清理因程序异常退出而遗留的 Ray worker 进程

echo "========================================="
echo "清理 Ray Worker 进程"
echo "========================================="

# 统计当前 Ray worker 进程数量
worker_count=$(pgrep -f "ray::TaskRunner" | wc -l)
echo "发现 $worker_count 个 Ray TaskRunner 进程"

if [ "$worker_count" -eq 0 ]; then
    echo "✓ 没有发现遗留的 Ray worker 进程"
    exit 0
fi

# 显示前10个进程信息
echo ""
echo "前10个进程信息:"
pgrep -f "ray::TaskRunner" | head -10 | xargs -I {} ps -p {} -o pid,cmd --no-headers | head -10

# 询问是否继续
echo ""
read -p "是否要终止所有这些进程? (y/N): " -n 1 -r
echo ""

if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "操作已取消"
    exit 0
fi

# 终止所有 Ray TaskRunner 进程
echo "正在终止进程..."
pkill -9 -f "ray::TaskRunner"

# 等待一小段时间让进程终止
sleep 1

# 检查是否还有残留进程
remaining=$(pgrep -f "ray::TaskRunner" | wc -l)
if [ "$remaining" -eq 0 ]; then
    echo "✓ 所有 Ray worker 进程已成功清理"
else
    echo "⚠ 警告: 仍有 $remaining 个进程未清理，可能需要手动检查"
    echo "剩余进程 PID:"
    pgrep -f "ray::TaskRunner"
fi

# 尝试清理其他 Ray 相关进程（可选）
echo ""
echo "检查其他 Ray 相关进程..."
raylet_count=$(pgrep -f "raylet" | wc -l)
gcs_count=$(pgrep -f "gcs_server" | wc -l)

if [ "$raylet_count" -gt 0 ] || [ "$gcs_count" -gt 0 ]; then
    echo "发现其他 Ray 进程:"
    [ "$raylet_count" -gt 0 ] && echo "  - $raylet_count 个 raylet 进程"
    [ "$gcs_count" -gt 0 ] && echo "  - $gcs_count 个 gcs_server 进程"
    echo "提示: 如果这些进程也属于已退出的 Ray 集群，可以使用以下命令清理:"
    echo "  pkill -9 -f raylet"
    echo "  pkill -9 -f gcs_server"
fi

echo ""
echo "========================================="
echo "清理完成"
echo "========================================="

